__author__ = 'Bill Liu'

from fractions import Fraction
import numpy as np
import copy

def matrix_parse(mat):
    # mat should be string
    mat = mat.strip()
    l = len(mat)
    val = []
    n_col = -1
    if l == 0 or mat[0] != '[' or mat[-1] != ']':
        print('Input format is wrong')
        return 0, 0, []
    for i in range(l):
        if mat[i] == '[':
            c = 0
            r = 0
            start = 0
            row = []
            continue
        if mat[i] == ' ' or mat[i] == ',':
            if start != 0:
                end = i
                num = Fraction(mat[start:end])
                row.append(num)
                start = 0
                c += 1
            continue
        if mat[i] == ';':
            if start != 0:
                end = i
                num = Fraction(mat[start:end])
                row.append(num)
                start = 0
                c += 1
            if n_col == -1:
                n_col = c
            else:
                assert(n_col == c)

            val.append(row)
            row = []
            r += 1
            c = 0
            continue
        if mat[i] == ']':
            if start != 0:
                end = i
                num = Fraction(mat[start:end])
                row.append(num)
                c += 1
                val.append(row)
                if n_col == -1:
                    n_col = c
                else:
                    assert(n_col == c)
            if c != 0:
                r += 1
            break

        if mat[i] != '/' and mat[i] != '-' and not mat[i].isdigit():
            print('Unrecognized input')
            return 0, 0, []
        else:
            if start == 0:
                start = i

    n_row = r
    return  val, n_row, n_col

def eye(n):
    mat = Matrix()
    mat.eye(n)
    return mat

def zeros(r, c):
    mat = Matrix()
    mat.zeros(r, c)
    return mat

def ones(r, c):
    mat = Matrix()
    mat.ones(r, c)
    return mat

def dot(*args):
    n_mat = len(args)
    mat = copy.deepcopy(args[0])
    for i in range(n_mat-1):
        mat.dot(args[i + 1])
    return mat

def divide(mat1, mat2):
    mat = copy.deepcopy(mat1)
    mat.divide(mat2)
    return mat


class Matrix:
    def __init__(self, mat=None):
        if mat is None:
            self.__data = []
            self.__row = 0
            self.__col = 0
        else:
            self.__data, self.__row, self.__col = matrix_parse(mat)

    def eye(self, n):
        self.__row = n
        self.__col = n
        self.__data = []
        for i in range(n):
            row = []
            for j in range(n):
                d = 1 if i == j else 0
                row.append(d)
            self.__data.append(row)

    def zeros(self, r, c):
        self.__row = r
        self.__col = c
        self.__data = []
        for i in range(r):
            row = []
            for j in range(c):
                row.append(0)
            self.__data.append(row)

    def ones(self, r, c):
        self.__row = r
        self.__col = c
        self.__data = []
        for i in range(r):
            row = []
            for j in range(c):
                row.append(1)
            self.__data.append(row)

    def set_element(self, row=None, col=None, val=None):
        if row is None or col is None or val is None:
            print('Invalid input parameter')
        self.__data[row][col] = val

    def get_element(self, r, c):
        return self.__data[r][c]

    def get_row_num(self):
        return self.__row

    def get_col_num(self):
        return self.__col

    def get_data(self):
        return copy.deepcopy(self.__data)

    def set_mat(self, data):
        self.__data = copy.deepcopy(data)
        self.__row = len(self.__data)
        self.__col = len(self.__data[0])

    def vstack(self, *mat2):
        num_mat = len(mat2)
        for i in range(num_mat):
            row2 = mat2[i].get_row_num()
            col2 = mat2[i].get_col_num()
            data2 = mat2[i].get_data()

            if col2 != self.__col:
                print('Oops, the columns of the input matrix do not match the original one')
            else:
                self.__row += row2
                self.__data += data2

    def hstack(self, *mat2):
        num_mat = len(mat2)
        for i in range(num_mat):
            row2 = mat2[i].get_row_num()
            col2 = mat2[i].get_col_num()
            data2 = mat2[i].get_data()

            if row2 != self.__row:
                print('Oops, the rows of the input matrix do not match the original one')
            else:
                self.__col += col2
                for r in range(row2):
                    self.__data[r] += data2[r]

    def get_col(self, col):
        if col >= self.__col:
            print('Column exceeds')
        mat = Matrix()
        mat.set_mat([[self.__data[i][col]] for i in range(self.__row)])
        return mat

    def get_row(self, row):
        if row >= self.__row:
            print('Row exceeds')
        mat = Matrix()
        mat.set_mat([self.__data[row]])
        return mat

    def get_col_slice(self, cols):
        data = []
        for i in range(self.__row):
            row = []
            for c in cols:
                row.append(self.__data[i][c])
            data.append(row)
        mat = Matrix()
        mat.set_mat(data)
        return mat

    def get_row_slice(self, rows):
        data = []
        for r in rows:
            if r >= self.__row:
                print('Row exceeds')
            data.append(self.__data[r])
        mat = Matrix()
        mat.set_mat(data)
        return mat

    def find_identity(self):
        n_var = self.__col
        n_rank = self.__row
        identity = np.zeros((n_rank,1))
        basis = []
        basis_pos = []
        for i in range(n_var):
            column = self.get_col(i).get_data()
            column_select = True
            pos_one = -1
            for j in range(n_rank):
                if column[j][0] != 0 and column[j][0] != 1:
                    column_select = False
                    break
                else:
                    if column[j][0] == 1:
                        if pos_one == -1:
                            pos_one = j
                        else:
                            column_select = False
                            break

            if pos_one == -1:
                column_select = False

            if column_select:
                basis.append(i)
                basis_pos.append(pos_one)
                identity[pos_one] = 1

        return basis, basis_pos

    def element_positive(self):
        is_positive = True
        for i in range(self.__row):
            for j in range(self.__col):
                if self.__data[i][j] < 0:
                    is_positive = False
                    break
            if not is_positive:
                break
        return is_positive

    def element_negative(self):
        is_negative = True
        for i in range(self.__row):
            for j in range(self.__col):
                if self.__data[i][j] > 0:
                    is_negative = False
                    break
            if not is_negative:
                break
        return is_negative

    def reformat(self, order):
        for i in range(self.__row):
            if i not in order:
                print('Some rows will be discarded after reformat')
        new_matrix = [self.__data[i] for i in order]
        self.__data = new_matrix
        self.__row = len(self.__data)
        return 0

    def get_col_width(self, col):
        width = 0
        for r in range(self.__row):
            coef = '{0}'.format(self.__data[r][col])
            width = max(width, len(coef))
        return width

    def format_print(self, bracket=True, hold_on=False, selected=None):
        col_width = np.zeros((self.__col,), dtype=int)

        for r in range(self.__row):
            for c in range(self.__col):
                coef = '{0}'.format(self.__data[r][c])
                col_width[c] = max(col_width[c], len(coef))
        total_width = int(np.sum(col_width) + self.__col * 2)

        println = []

        if bracket is True:
            line = ' _'
            for i in range(total_width - 2):
                line += ' '
            line += '_ '
        else:
            line = ''
            for i in range(total_width + 2):
                line += ' '
        println.append(line)

        for r in range(self.__row):
            line = ' '
            if bracket is True:
                line = '|'
            for c in range(self.__col):
                coef = '{0}'.format(self.__data[r][c])
                line += ' '
                for i in range(col_width[c]-len(coef)):
                    line += ' '
                line += coef
                line += ' '
            if bracket is True:
                line += '|'
            else:
                line += ' '
            println.append(line)

            if r != self.__row - 1:
                line = ' '
                if bracket is True:
                    line = '|'
                for i in range(total_width):
                    line += ' '
                if bracket is True:
                    line += '|'
                else:
                    line += ' '
                println.append(line)

        if bracket is True:
            line = ' -'
            for i in range(total_width - 2):
                line += ' '
            line += '- '
        else:
            line = ''
            for i in range(total_width):
                line += ' '
        println.append(line)

        if selected is not None:
            sr = selected[0]
            sc = selected[1]

            if sr > self.__row or sc > self.__col:
                print('Selected element exceeds index!')

            base_row = sr * 2
            base_col = 1 + sc * 2 + int(np.sum(col_width[:sc]))

            i = 1
            println[base_row + i] = println[base_row + i][:base_col] + '>' + println[base_row + i][base_col+1:]
            println[base_row + i] = println[base_row + i][:base_col+col_width[sc]+1] + '<' + println[base_row + i][base_col+col_width[sc]+2:]

            i = 2
            line = println[base_row + i][:base_col+1]
            for j in range(col_width[sc]):
                line += '^'
            println[base_row + i] = line + println[base_row + i][base_col + col_width[sc]+1:]

        if hold_on is False:
            for line in println:
                print(line)

        return println

    def dot(self, mat2):
        n_row2 = mat2.get_row_num()
        n_col2 = mat2.get_col_num()

        data = []

        mat2 = mat2.get_data()
        assert(self.__col == n_row2)
        for i in range(self.__row):
            row = []
            for j in range(n_col2):
                prod = 0
                for k in range(self.__col):
                    prod += self.__data[i][k] * mat2[k][j]
                row.append(prod)
            data.append(row)

        self.__col = n_col2
        self.__data = data
        return self

    def divide(self, mat2):
        assert(self.__col == mat2.get_col_num() and self.__row == mat2.get_row_num())
        sub = mat2.get_data()
        for i in range(self.__row):
            for j in range(self.__col):
                if sub[i][j] < 1e-10:
                    self.__data[i][j] = 1e10
                else:
                    self.__data[i][j] /= sub[i][j]
        return self

    def subtract(self, mat2):
        assert(self.__col == mat2.get_col_num() and self.__row == mat2.get_row_num())
        sub = mat2.get_data()
        for i in range(self.__row):
            for j in range(self.__col):
                self.__data[i][j] -= sub[i][j]
        return self

    def gaussian_elimination(self, row, col):
        factor = self.__data[row][col]
        # set this element to 1
        self.__data[row] = [self.__data[row][i] / factor for i in range(self.__col)]
        # eliminate other elements in that column
        for i in range(self.__row):
            if i == row:
                continue

            factor = self.__data[i][col]
            self.__data[i] = [(self.__data[i][j] - self.__data[row][j] * factor) for j in range(self.__col)]

